import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
from scipy.optimize import minimize
from datetime import datetime, timedelta
import warnings

warnings.filterwarnings('ignore')

# 设置matplotlib显示中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def load_stock_data():
    """
    加载当前目录或上一级目录下的所有股票数据CSV文件
    假设每个文件名为股票代码命名，如 000001.SZ.csv
    """
    stock_data = {}
    valid_stocks = []

    # 获取当前工作目录
    current_dir = os.getcwd()
    print(f"当前工作目录: {current_dir}")

    # 查找当前目录下的所有CSV文件
    csv_files = [f for f in os.listdir(current_dir) if f.endswith('.csv') and f != '中国平安2318未来一个月预测.csv']

    # 如果当前目录没有找到CSV文件，则查找上一级目录
    if len(csv_files) == 0:
        print("当前目录未找到CSV文件，尝试在上一级目录查找...")
        parent_dir = os.path.dirname(current_dir)
        csv_files = [f for f in os.listdir(parent_dir) if f.endswith('.csv') and f != '中国平安2318未来一个月预测.csv']
        if len(csv_files) > 0:
            print(f"在上一级目录 {parent_dir} 中找到 {len(csv_files)} 个CSV文件")
            current_dir = parent_dir
        else:
            print("警告: 上一级目录中也没有找到CSV股票数据文件")

    if len(csv_files) == 0:
        print("警告: 当前目录下没有找到CSV股票数据文件")
        # 尝试在常见的数据目录中查找
        possible_dirs = [
            os.path.join(current_dir, 'data'),
            os.path.join(current_dir, 'stock_data'),
            current_dir
        ]

        for data_dir in possible_dirs:
            if os.path.exists(data_dir):
                csv_files = [f for f in os.listdir(data_dir) if
                             f.endswith('.csv') and f != '中国平安2318未来一个月预测.csv']
                if len(csv_files) > 0:
                    print(f"在目录 {data_dir} 中找到 {len(csv_files)} 个CSV文件")
                    current_dir = data_dir
                    break

    if len(csv_files) == 0:
        print("错误: 没有找到任何股票数据CSV文件")
        return stock_data, valid_stocks

    for filename in csv_files:
        stock_code = filename.replace('.csv', '')
        file_path = os.path.join(current_dir, filename)
        try:
            df = pd.read_csv(file_path)
            # 确保必要的列存在
            required_columns = ['日期', '开盘', '高', '低', '收盘', '交易量']
            if not all(col in df.columns for col in required_columns):
                print(f"警告: {filename} 缺少必要列，跳过")
                continue

            # 转换日期列为 datetime 类型
            df['日期'] = pd.to_datetime(df['日期'])
            df = df.sort_values('日期').reset_index(drop=True)

            stock_data[stock_code] = df
            valid_stocks.append(stock_code)
            print(f"成功加载 {filename}")
        except Exception as e:
            print(f"警告: 无法加载 {filename}，错误: {e}")

    print(f"成功加载 {len(valid_stocks)} 支股票数据: {valid_stocks}")
    return stock_data, valid_stocks


def preprocess_data(stock_data, valid_stocks):
    """预处理数据，计算收益率等"""
    print("\n===== 数据预处理 =====")

    # 统一日期范围
    all_dates = []
    for stock in valid_stocks:
        all_dates.extend(stock_data[stock]['日期'].tolist())

    all_dates = sorted(list(set(all_dates)))

    # 检查是否有足够的日期数据
    if len(all_dates) == 0:
        print("错误: 没有可用的日期数据")
        return None, None, []

    # 降低共同日期的要求，从70%降到至少2支股票
    common_dates = [date for date in all_dates if all_dates.count(date) >= min(2, len(valid_stocks))]

    # 如果仍然没有共同日期，尝试更宽松的条件
    if len(common_dates) == 0:
        print("警告: 没有找到足够的共同交易日，使用所有日期")
        common_dates = all_dates

    if len(common_dates) == 0:
        print("错误: 无法确定交易日期范围")
        return None, None, []

    print(f"确定共同交易日范围: {common_dates[0].date()} 至 {common_dates[-1].date()}，共 {len(common_dates)} 天")

    # 创建收益率矩阵，使用common_dates作为索引
    returns_matrix = pd.DataFrame(index=common_dates)

    # 检测数据异常（如完全相同的数据）
    suspicious_pairs = []
    for i in range(len(valid_stocks)):
        for j in range(i + 1, len(valid_stocks)):
            stock1, stock2 = valid_stocks[i], valid_stocks[j]

            # 检查收盘价是否完全相同
            df1 = stock_data[stock1][stock_data[stock1]['日期'].isin(common_dates)].sort_values('日期')
            df2 = stock_data[stock2][stock_data[stock2]['日期'].isin(common_dates)].sort_values('日期')

            if len(df1) == len(df2) and len(df1) > 0 and np.allclose(df1['收盘'].values, df2['收盘'].values, atol=0.01):
                suspicious_pairs.append((stock1, stock2))

    if suspicious_pairs:
        print("\n检测到可疑的重复数据:")
        for stock1, stock2 in suspicious_pairs:
            print(f"  {stock1} 和 {stock2} 的价格数据高度相似，可能存在数据错误")

        # 从可疑对中移除一个股票（保留第一个）
        stocks_to_remove = set()
        for stock1, stock2 in suspicious_pairs:
            stocks_to_remove.add(stock2)  # 移除第二个

        valid_stocks = [stock for stock in valid_stocks if stock not in stocks_to_remove]
        print(f"\n移除重复股票后，剩余 {len(valid_stocks)} 支有效股票: {valid_stocks}")

        # 重新计算共同日期
        all_dates = []
        for stock in valid_stocks:
            all_dates.extend(stock_data[stock]['日期'].tolist())
        all_dates = sorted(list(set(all_dates)))
        common_dates = [date for date in all_dates if all_dates.count(date) >= min(2, len(valid_stocks))]

        if len(common_dates) == 0:
            common_dates = all_dates

        # 重新创建收益率矩阵
        returns_matrix = pd.DataFrame(index=common_dates)

    # 计算每只股票的收益率，并确保与returns_matrix索引对齐
    for stock in valid_stocks:
        df = stock_data[stock].copy()
        # 筛选共同日期的数据并排序
        df = df[df['日期'].isin(common_dates)].sort_values('日期').reset_index(drop=True)

        # 计算日收益率
        df['收益率'] = df['收盘'].pct_change()

        # 填充缺失值（如果有的话）
        df['收益率'] = df['收益率'].fillna(0)

        # 创建一个与returns_matrix索引对齐的序列
        aligned_returns = pd.Series(df['收益率'].values, index=df['日期'])
        aligned_returns = aligned_returns.reindex(returns_matrix.index, fill_value=0)

        # 将收益率添加到矩阵
        returns_matrix[stock] = aligned_returns.values

    # 检查协方差矩阵是否有效
    cov_matrix = returns_matrix.cov()
    # 修复：正确检查协方差矩阵中是否包含NaN或Inf值
    if cov_matrix.isna().any().any() or np.isinf(cov_matrix).any().any():
        print("警告: 协方差矩阵包含NaN或Inf值，进行数据清理...")
        # 移除有异常值的股票
        valid_stocks = [stock for stock in valid_stocks if not returns_matrix[stock].isna().any()]
        returns_matrix = returns_matrix[valid_stocks]
        cov_matrix = returns_matrix.cov()

    print(f"\n最终收益率矩阵维度: {returns_matrix.shape}")
    return returns_matrix, cov_matrix, valid_stocks

def calculate_scheme_prices(df):
    """计算两种交易方案的价格"""
    # 方案1: 买入价格=最低价，卖出价格=最高价
    buy_price1 = df['低'].values
    sell_price1 = df['高'].values

    # 方案2: 买入价格=最低价+0.25*(最高价-最低价)，卖出价格=最低价+0.75*(最高价-最低价)
    price_range = df['高'].values - df['低'].values
    buy_price2 = df['低'].values + price_range * 0.25
    sell_price2 = df['低'].values + price_range * 0.75

    return buy_price1, sell_price1, buy_price2, sell_price2


def adjust_returns_for_trading_scheme(stock_data, valid_stocks, returns_matrix, transaction_fee=0.0003):
    """
    根据交易方案调整收益率
    考虑T+1交易规则和交易费用
    """
    print("\n===== 调整收益率以适应交易方案 =====")

    # 为每种方案创建调整后的收益率矩阵
    adjusted_returns1 = returns_matrix.copy()  # 方案1
    adjusted_returns2 = returns_matrix.copy()  # 方案2

    for stock in valid_stocks:
        df = stock_data[stock]

        # 确保数据按日期排序
        df = df.sort_values('日期')

        # 计算两种方案的实际交易价格
        buy_price1, sell_price1, buy_price2, sell_price2 = calculate_scheme_prices(df)

        # 计算调整后的收益率
        for i in range(1, len(df)):
            # 获取日期
            current_date = df.iloc[i]['日期']

            # 检查当前日期是否在共同交易日中
            if current_date not in returns_matrix.index:
                continue

            # 方案1: 买入=最低价，卖出=最高价
            # 实际收益率 = (卖出价格*(1-交易费) - 买入价格*(1+交易费)) / 买入价格*(1+交易费)
            # 由于T+1规则，今天买入的价格是今天的最低价，明天卖出的价格是明天的最高价
            if i < len(df) - 1:
                buy_price = buy_price1[i]
                sell_price = sell_price1[i + 1]

                # 考虑交易费用
                cost = buy_price * (1 + transaction_fee)
                revenue = sell_price * (1 - transaction_fee)
                actual_return = (revenue - cost) / cost

                # 更新调整后的收益率
                idx = returns_matrix.index.get_loc(current_date)
                if idx < len(adjusted_returns1) - 1:
                    adjusted_returns1.at[current_date, stock] = actual_return

            # 方案2: 买入=1/4分位，卖出=3/4分位
            if i < len(df) - 1:
                buy_price = buy_price2[i]
                sell_price = sell_price2[i + 1]

                # 考虑交易费用
                cost = buy_price * (1 + transaction_fee)
                revenue = sell_price * (1 - transaction_fee)
                actual_return = (revenue - cost) / cost

                # 更新调整后的收益率
                idx = returns_matrix.index.get_loc(current_date)
                if idx < len(adjusted_returns2) - 1:
                    adjusted_returns2.at[current_date, stock] = actual_return

    # 计算调整后的年化收益率和协方差矩阵
    annualized_returns1 = adjusted_returns1.mean() * 252
    cov_matrix1 = adjusted_returns1.cov() * 252

    annualized_returns2 = adjusted_returns2.mean() * 252
    cov_matrix2 = adjusted_returns2.cov() * 252

    print(f"方案1: 年化收益率范围: {annualized_returns1.min():.4f} 到 {annualized_returns1.max():.4f}")
    print(f"方案2: 年化收益率范围: {annualized_returns2.min():.4f} 到 {annualized_returns2.max():.4f}")

    return {
        'scheme1': {
            'returns': adjusted_returns1,
            'annualized_returns': annualized_returns1,
            'cov_matrix': cov_matrix1
        },
        'scheme2': {
            'returns': adjusted_returns2,
            'annualized_returns': annualized_returns2,
            'cov_matrix': cov_matrix2
        }
    }


def optimize_portfolio(annualized_returns, cov_matrix, risk_free_rate=0.02):
    """
    优化投资组合
    使用最小化负夏普比率的方法
    """
    num_assets = len(annualized_returns)

    # 约束条件：权重和为1
    constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})

    # 边界条件：权重在0到1之间（不允许做空）
    bounds = tuple((0, 1) for _ in range(num_assets))

    # 初始权重（等权重）
    init_weights = num_assets * [1. / num_assets]

    # 目标函数：最大化夏普比率 = (组合收益率 - 无风险利率) / 组合风险
    # 由于优化器是求最小值，所以取负号
    def neg_sharpe_ratio(weights):
        portfolio_return = np.sum(annualized_returns * weights)
        portfolio_volatility = np.sqrt(np.dot(weights.T, np.dot(cov_matrix, weights)))
        sharpe_ratio = (portfolio_return - risk_free_rate) / portfolio_volatility
        return -sharpe_ratio

    # 执行优化
    result = minimize(neg_sharpe_ratio, init_weights, method='SLSQP', bounds=bounds, constraints=constraints)

    # 提取最优权重
    optimal_weights = result.x

    # 计算最优组合的指标
    portfolio_return = np.sum(annualized_returns * optimal_weights)
    portfolio_volatility = np.sqrt(np.dot(optimal_weights.T, np.dot(cov_matrix, optimal_weights)))
    sharpe_ratio = (portfolio_return - risk_free_rate) / portfolio_volatility

    return optimal_weights, portfolio_return, portfolio_volatility, sharpe_ratio


def generate_efficient_frontier(annualized_returns, cov_matrix, risk_free_rate=0.02, num_points=50):
    """
    生成有效前沿
    """
    num_assets = len(annualized_returns)

    # 约束条件：权重和为1
    constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})

    # 边界条件：权重在0到1之间（不允许做空）
    bounds = tuple((0, 1) for _ in range(num_assets))

    # 计算最小和最大可能的收益率
    min_return = annualized_returns.min()
    max_return = annualized_returns.max()

    # 生成目标收益率序列
    target_returns = np.linspace(min_return, max_return, num_points)

    # 存储结果
    portfolio_returns = []
    portfolio_risks = []
    portfolio_weights = []

    for target_return in target_returns:
        # 目标函数：最小化组合风险（方差）
        def portfolio_variance(weights):
            return np.dot(weights.T, np.dot(cov_matrix, weights))

        # 约束条件：权重和为1，且组合收益率等于目标收益率
        constraints = (
            {'type': 'eq', 'fun': lambda x: np.sum(x) - 1},
            {'type': 'eq', 'fun': lambda x: np.sum(annualized_returns * x) - target_return}
        )

        # 执行优化
        result = minimize(portfolio_variance, num_assets * [1. / num_assets],
                          method='SLSQP', bounds=bounds, constraints=constraints)

        if result.success:
            weights = result.x
            volatility = np.sqrt(result.fun)

            portfolio_returns.append(target_return)
            portfolio_risks.append(volatility)
            portfolio_weights.append(weights)

    # 检查是否有有效的投资组合
    if len(portfolio_returns) == 0:
        print("警告: 未能生成有效前沿，所有优化都失败了")
        # 返回默认值
        return {
            'returns': [],
            'risks': [],
            'weights': [],
            'tangent': {
                'return': 0,
                'risk': 1,
                'weights': np.array([1.0 / num_assets] * num_assets) if num_assets > 0 else np.array([]),
                'sharpe': 0
            }
        }

    # 计算切线组合（夏普比率最高的组合）
    # 使用try-except处理可能的异常
    try:
        sharpe_ratios = [(portfolio_returns[i] - risk_free_rate) / portfolio_risks[i]
                         for i in range(len(portfolio_returns)) if portfolio_risks[i] > 0]

        if len(sharpe_ratios) > 0:
            max_sharpe_idx = np.argmax(sharpe_ratios)
            # 需要找出原始索引
            valid_indices = [i for i in range(len(portfolio_returns)) if portfolio_risks[i] > 0]
            original_max_sharpe_idx = valid_indices[max_sharpe_idx]

            tangent_return = portfolio_returns[original_max_sharpe_idx]
            tangent_risk = portfolio_risks[original_max_sharpe_idx]
            tangent_weights = portfolio_weights[original_max_sharpe_idx]
            tangent_sharpe = (tangent_return - risk_free_rate) / tangent_risk
        else:
            # 如果所有风险都为0，则选择收益最高的组合
            max_return_idx = np.argmax(portfolio_returns)
            tangent_return = portfolio_returns[max_return_idx]
            tangent_risk = portfolio_risks[max_return_idx]
            tangent_weights = portfolio_weights[max_return_idx]
            tangent_sharpe = 0 if tangent_risk == 0 else (tangent_return - risk_free_rate) / tangent_risk

    except Exception as e:
        print(f"警告: 计算切线组合时出错: {e}")
        # 默认使用第一个有效的组合
        tangent_return = portfolio_returns[0]
        tangent_risk = portfolio_risks[0]
        tangent_weights = portfolio_weights[0]
        tangent_sharpe = (tangent_return - risk_free_rate) / tangent_risk if tangent_risk > 0 else 0

    return {
        'returns': portfolio_returns,
        'risks': portfolio_risks,
        'weights': portfolio_weights,
        'tangent': {
            'return': tangent_return,
            'risk': tangent_risk,
            'weights': tangent_weights,
            'sharpe': tangent_sharpe
        }
    }


def backtest_portfolio(stock_data, optimal_weights, valid_stocks, start_date, end_date, initial_capital=1000000,
                       transaction_fee=0.0003):
    """
    回测投资组合
    """
    print("\n===== 投资组合回测 =====")
    print(f"回测周期: {start_date.date()} 至 {end_date.date()}")
    print(f"初始资金: {initial_capital:,.2f}元")

    # 筛选回测期间的数据
    portfolio_value = []
    dates = []
    daily_returns = []

    # 获取所有股票在回测期间的数据
    stock_dfs = {}
    for stock in valid_stocks:
        df = stock_data[stock].copy()
        df = df[(df['日期'] >= start_date) & (df['日期'] <= end_date)].sort_values('日期').reset_index(drop=True)
        if len(df) > 0:
            stock_dfs[stock] = df

    # 确定共同的交易日
    common_dates = None
    for stock, df in stock_dfs.items():
        if common_dates is None:
            common_dates = set(df['日期'].tolist())
        else:
            common_dates = common_dates & set(df['日期'].tolist())

    common_dates = sorted(list(common_dates))

    if not common_dates:
        print("错误: 没有共同的交易日，无法进行回测")
        return None

    print(f"回测期间共有 {len(common_dates)} 个交易日")

    # 初始化投资组合
    capital = initial_capital
    holdings = {stock: 0 for stock in valid_stocks}

    # 按日期进行回测
    for i, date in enumerate(common_dates):
        # 当天的总资产价值
        total_value = capital
        for stock in valid_stocks:
            if stock in stock_dfs and date in stock_dfs[stock]['日期'].values:
                df = stock_dfs[stock]
                idx = df[df['日期'] == date].index[0]
                price = df.iloc[idx]['收盘']
                total_value += holdings[stock] * price

        portfolio_value.append(total_value)
        dates.append(date)

        # 如果不是最后一天，计算当天的收益率
        if i > 0:
            daily_return = (portfolio_value[-1] / portfolio_value[-2]) - 1
            daily_returns.append(daily_return)

        # 如果是第一个交易日，初始化投资组合
        if i == 0:
            for j, stock in enumerate(valid_stocks):
                if optimal_weights[j] > 0 and stock in stock_dfs and date in stock_dfs[stock]['日期'].values:
                    df = stock_dfs[stock]
                    idx = df[df['日期'] == date].index[0]
                    price = df.iloc[idx]['收盘']

                    # 计算可购买的数量（考虑交易费用）
                    investment = initial_capital * optimal_weights[j]
                    shares = investment / (price * (1 + transaction_fee))

                    # 更新持仓和现金
                    holdings[stock] = shares
                    capital -= shares * price * (1 + transaction_fee)

    # 计算性能指标
    portfolio_returns = np.array(daily_returns)
    annualized_return = (portfolio_value[-1] / portfolio_value[0]) ** (252 / len(common_dates)) - 1
    annualized_volatility = np.std(portfolio_returns) * np.sqrt(252)
    sharpe_ratio = (annualized_return - 0.02) / annualized_volatility if annualized_volatility > 0 else 0

    print(f"回测结果:")
    print(f"  最终资产: {portfolio_value[-1]:,.2f}元")
    print(f"  总收益率: {portfolio_value[-1] / initial_capital - 1:.4%}")
    print(f"  年化收益率: {annualized_return:.4%}")
    print(f"  年化波动率: {annualized_volatility:.4%}")
    print(f"  夏普比率: {sharpe_ratio:.4f}")

    # 创建回测结果DataFrame
    backtest_results = pd.DataFrame({
        '日期': dates,
        '投资组合价值': portfolio_value
    })

    return {
        'backtest_results': backtest_results,
        'annualized_return': annualized_return,
        'annualized_volatility': annualized_volatility,
        'sharpe_ratio': sharpe_ratio,
        'final_value': portfolio_value[-1]
    }


def plot_results(adjusted_data, valid_stocks, efficient_frontier1, efficient_frontier2,
                 optimal_weights1, optimal_weights2, backtest_results1, backtest_results2):
    """绘制结果图表"""
    print("\n===== 生成结果图表 =====")

    # 1. 股票收益率和风险散点图
    plt.figure(figsize=(14, 7))

    # 提取每只股票的年化收益率和波动率
    returns1 = adjusted_data['scheme1']['annualized_returns']
    volatilities1 = np.sqrt(np.diag(adjusted_data['scheme1']['cov_matrix']))

    returns2 = adjusted_data['scheme2']['annualized_returns']
    volatilities2 = np.sqrt(np.diag(adjusted_data['scheme2']['cov_matrix']))

    # 绘制方案1的散点图
    plt.scatter(volatilities1, returns1, c='blue', alpha=0.6, s=100, label='方案1: 买入=最低价，卖出=最高价')

    # 绘制方案2的散点图
    plt.scatter(volatilities2, returns2, c='red', alpha=0.6, s=100, label='方案2: 买入=1/4分位，卖出=3/4分位')

    # 标记股票名称
    for i, stock in enumerate(valid_stocks):
        plt.annotate(stock, (volatilities1[i] * 1.02, returns1[i] * 1.02), fontsize=9)

    plt.title('各股票收益率-风险分析', fontsize=16)
    plt.xlabel('年化波动率', fontsize=12)
    plt.ylabel('年化收益率', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig('stock_risk_return.png', dpi=300, bbox_inches='tight')
    print("已保存股票收益率-风险分析图: stock_risk_return.png")

    # 2. 有效前沿对比
    plt.figure(figsize=(14, 7))

    # 绘制方案1的有效前沿
    if len(efficient_frontier1['risks']) > 0:
        plt.plot(efficient_frontier1['risks'], efficient_frontier1['returns'], 'b-', label='方案1有效前沿')
        plt.scatter(efficient_frontier1['tangent']['risk'], efficient_frontier1['tangent']['return'],
                    c='blue', s=100, marker='*', label='方案1切线组合')

    # 绘制方案2的有效前沿
    if len(efficient_frontier2['risks']) > 0:
        plt.plot(efficient_frontier2['risks'], efficient_frontier2['returns'], 'r-', label='方案2有效前沿')
        plt.scatter(efficient_frontier2['tangent']['risk'], efficient_frontier2['tangent']['return'],
                    c='red', s=100, marker='*', label='方案2切线组合')

    # 添加无风险利率线
    risk_free_rate = 0.02
    # 安全地计算最大风险值
    all_risks = []
    if len(efficient_frontier1['risks']) > 0:
        all_risks.extend(efficient_frontier1['risks'])
    if len(efficient_frontier2['risks']) > 0:
        all_risks.extend(efficient_frontier2['risks'])

    if len(all_risks) > 0:
        max_risk = max(all_risks)
        plt.plot([0, max_risk], [risk_free_rate, risk_free_rate + max_risk * 1.5],
                 'g--', label='资本市场线')
    else:
        # 如果没有有效前沿数据，绘制一个默认的资本市场线
        plt.plot([0, 1], [risk_free_rate, risk_free_rate + 1 * 1.5],
                 'g--', label='资本市场线')

    plt.title('投资组合有效前沿对比', fontsize=16)
    plt.xlabel('年化波动率', fontsize=12)
    plt.ylabel('年化收益率', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig('efficient_frontier.png', dpi=300, bbox_inches='tight')
    print("已保存有效前沿对比图: efficient_frontier.png")

    # 3. 最优权重对比
    plt.figure(figsize=(14, 7))

    # 准备数据
    x = np.arange(len(valid_stocks))
    width = 0.35

    # 绘制方案1的权重
    plt.bar(x - width / 2, optimal_weights1, width, label='方案1: 买入=最低价，卖出=最高价', color='blue', alpha=0.7)

    # 绘制方案2的权重
    plt.bar(x + width / 2, optimal_weights2, width, label='方案2: 买入=1/4分位，卖出=3/4分位', color='red', alpha=0.7)

    # 添加标签和标题
    plt.xlabel('股票名称', fontsize=12)
    plt.ylabel('投资权重', fontsize=12)
    plt.title('最优投资组合权重对比', fontsize=16)
    plt.xticks(x, valid_stocks, rotation=45, ha='right')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig('portfolio_weights.png', dpi=300, bbox_inches='tight')
    print("已保存投资组合权重对比图: portfolio_weights.png")

    # 4. 投资组合价值对比
    plt.figure(figsize=(14, 7))

    # 绘制方案1的组合价值
    if backtest_results1 is not None and 'backtest_results' in backtest_results1:
        plt.plot(backtest_results1['backtest_results']['日期'],
                 backtest_results1['backtest_results']['投资组合价值'],
                 label=f"方案1: 年化收益 {backtest_results1['annualized_return']:.2%}, 夏普 {backtest_results1['sharpe_ratio']:.2f}",
                 linewidth=2.5)

    # 绘制方案2的组合价值
    if backtest_results2 is not None and 'backtest_results' in backtest_results2:
        plt.plot(backtest_results2['backtest_results']['日期'],
                 backtest_results2['backtest_results']['投资组合价值'],
                 label=f"方案2: 年化收益 {backtest_results2['annualized_return']:.2%}, 夏普 {backtest_results2['sharpe_ratio']:.2f}",
                 linewidth=2.5)

    # 添加初始资金线
    if (backtest_results1 is not None and 'backtest_results' in backtest_results1 and
            len(backtest_results1['backtest_results']) > 0):
        initial_capital = backtest_results1['backtest_results']['投资组合价值'].iloc[0]
        plt.axhline(y=initial_capital, color='r', linestyle='--', label='初始资金', linewidth=1.5)

    plt.title('投资组合价值对比', fontsize=16)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('投资组合价值(元)', fontsize=12)
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)

    # 设置日期格式
    plt.gca().xaxis.set_major_formatter(plt.matplotlib.dates.DateFormatter('%m-%d'))
    plt.gca().xaxis.set_major_locator(plt.matplotlib.dates.DayLocator(interval=30))
    plt.gcf().autofmt_xdate()

    plt.tight_layout()
    plt.savefig('portfolio_value.png', dpi=300, bbox_inches='tight')
    print("已保存投资组合价值对比图: portfolio_value.png")

    # 5. 组合构成表格
    plt.figure(figsize=(14, 8))
    plt.axis('off')

    # 创建表格数据
    table_data = []
    for i, stock in enumerate(valid_stocks):
        table_data.append([
            stock,
            f"{optimal_weights1[i]:.4f}",
            f"{optimal_weights2[i]:.4f}",
            f"{adjusted_data['scheme1']['annualized_returns'][stock]:.4f}",
            f"{adjusted_data['scheme2']['annualized_returns'][stock]:.4f}",
            f"{np.sqrt(adjusted_data['scheme1']['cov_matrix'].loc[stock, stock]):.4f}",
            f"{np.sqrt(adjusted_data['scheme2']['cov_matrix'].loc[stock, stock]):.4f}"
        ])

    # 创建表格
    columns = ['股票名称', '方案1权重', '方案2权重', '方案1年化收益', '方案2年化收益', '方案1年化波动', '方案2年化波动']
    table = plt.table(cellText=table_data,
                      colLabels=columns,
                      cellLoc='center',
                      loc='center',
                      bbox=[0.1, 0.1, 0.8, 0.8])

    table.auto_set_font_size(False)
    table.set_fontsize(9)
    table.scale(1, 1.5)

    plt.tight_layout()
    plt.savefig('portfolio_composition.png', dpi=300, bbox_inches='tight')
    print("已保存投资组合构成表格: portfolio_composition.png")

    plt.show()


def save_results(stock_data, valid_stocks, optimal_weights1, optimal_weights2,
                 efficient_frontier1, efficient_frontier2,
                 backtest_results1, backtest_results2, adjusted_data):
    """保存结果到文件"""
    print("\n===== 保存结果到文件 =====")

    # 1. 保存最优权重
    weights_df = pd.DataFrame({
        '股票名称': valid_stocks,
        '方案1权重': optimal_weights1,
        '方案2权重': optimal_weights2,
        '方案1年化收益': adjusted_data['scheme1']['annualized_returns'].values,
        '方案2年化收益': adjusted_data['scheme2']['annualized_returns'].values,
        '方案1年化波动': np.sqrt(np.diag(adjusted_data['scheme1']['cov_matrix'])),
        '方案2年化波动': np.sqrt(np.diag(adjusted_data['scheme2']['cov_matrix']))
    })
    weights_df.to_csv('optimal_weights.csv', index=False)
    print("已保存最优权重: optimal_weights.csv")

    # 2. 保存有效前沿数据
    if len(efficient_frontier1['risks']) > 0:
        frontier_data1 = pd.DataFrame({
            '风险': efficient_frontier1['risks'],
            '收益': efficient_frontier1['returns']
        })
        frontier_data1.to_csv('efficient_frontier_scheme1.csv', index=False)

    if len(efficient_frontier2['risks']) > 0:
        frontier_data2 = pd.DataFrame({
            '风险': efficient_frontier2['risks'],
            '收益': efficient_frontier2['returns']
        })
        frontier_data2.to_csv('efficient_frontier_scheme2.csv', index=False)
    print("已保存有效前沿数据")

    # 3. 保存回测结果
    if backtest_results1 is not None and 'backtest_results' in backtest_results1:
        backtest_results1['backtest_results'].to_csv('backtest_results_scheme1.csv', index=False)
    if backtest_results2 is not None and 'backtest_results' in backtest_results2:
        backtest_results2['backtest_results'].to_csv('backtest_results_scheme2.csv', index=False)
    print("已保存回测结果")

    # 4. 生成策略报告
    # 安全地获取切线组合的指标
    ef1_sharpe = efficient_frontier1['tangent']['sharpe'] if 'tangent' in efficient_frontier1 else 0
    ef2_sharpe = efficient_frontier2['tangent']['sharpe'] if 'tangent' in efficient_frontier2 else 0
    ef1_return = efficient_frontier1['tangent']['return'] if 'tangent' in efficient_frontier1 else 0
    ef2_return = efficient_frontier2['tangent']['return'] if 'tangent' in efficient_frontier2 else 0
    ef1_risk = efficient_frontier1['tangent']['risk'] if 'tangent' in efficient_frontier1 else 1
    ef2_risk = efficient_frontier2['tangent']['risk'] if 'tangent' in efficient_frontier2 else 1

    # 安全地获取回测结果
    br1_annual_return = backtest_results1['annualized_return'] if backtest_results1 is not None else 0
    br2_annual_return = backtest_results2['annualized_return'] if backtest_results2 is not None else 0
    br1_volatility = backtest_results1['annualized_volatility'] if backtest_results1 is not None else 0
    br2_volatility = backtest_results2['annualized_volatility'] if backtest_results2 is not None else 0
    br1_sharpe = backtest_results1['sharpe_ratio'] if backtest_results1 is not None else 0
    br2_sharpe = backtest_results2['sharpe_ratio'] if backtest_results2 is not None else 0

    report = f"""
# 多支股票组合投资策略报告

## 基本信息
- **股票池**: {', '.join(valid_stocks)}
- **数据时间范围**: {min([df['日期'].min() for df in stock_data.values()]).date()} 至 {max([df['日期'].max() for df in stock_data.values()]).date()}
- **无风险利率**: 2% (年化)
- **交易费用**: 0.03%

## 交易方案比较

### 方案1: 买入=最低价，卖出=最高价
- **最优投资组合年化收益率**: {ef1_return:.4%}
- **最优投资组合年化波动率**: {ef1_risk:.4%}
- **最优投资组合夏普比率**: {ef1_sharpe:.4f}
- **回测年化收益率**: {br1_annual_return:.4%}
- **回测年化波动率**: {br1_volatility:.4%}
- **回测夏普比率**: {br1_sharpe:.4f}

### 方案2: 买入=1/4分位，卖出=3/4分位
- **最优投资组合年化收益率**: {ef2_return:.4%}
- **最优投资组合年化波动率**: {ef2_risk:.4%}
- **最优投资组合夏普比率**: {ef2_sharpe:.4f}
- **回测年化收益率**: {br2_annual_return:.4%}
- **回测年化波动率**: {br2_volatility:.4%}
- **回测夏普比率**: {br2_sharpe:.4f}

## 最优投资组合构成

### 方案1: 买入=最低价，卖出=最高价
"""

    for i, stock in enumerate(valid_stocks):
        if optimal_weights1[i] > 0.01:  # 只显示权重超过1%的股票
            report += f"- {stock}: {optimal_weights1[i]:.2%}\n"

    report += "\n### 方案2: 买入=1/4分位，卖出=3/4分位\n"
    for i, stock in enumerate(valid_stocks):
        if optimal_weights2[i] > 0.01:  # 只显示权重超过1%的股票
            report += f"- {stock}: {optimal_weights2[i]:.2%}\n"

    report += f"""
## 策略结论
{'方案1' if ef1_sharpe > ef2_sharpe else '方案2'} 
的夏普比率更高，表明{'该方案' if ef1_sharpe > ef2_sharpe else '该方案'} 
在风险调整后的收益表现更好。

{'方案1的夏普比率比方案2高' + f"{ef1_sharpe - ef2_sharpe:.4f}"
    if ef1_sharpe > ef2_sharpe
    else '方案2的夏普比率比方案1高' + f"{ef2_sharpe - ef1_sharpe:.4f}" if ef2_sharpe > ef1_sharpe else '两个方案的夏普比率相同'}

## 建议
{'建议采用方案1进行投资组合构建，因为它提供了更高的风险调整后收益。'
    if ef1_sharpe > ef2_sharpe
    else '建议采用方案2进行投资组合构建，因为它提供了更高的风险调整后收益。' if ef2_sharpe > ef1_sharpe else '两个方案的风险调整后收益相同，可根据其他因素选择。'}

投资组合应重点关注:
"""

    # 找出两个方案中权重最高的3只股票
    top_stocks = set()
    if len(optimal_weights1) > 0:
        top_indices1 = np.argsort(optimal_weights1)[-3:]
        for idx in top_indices1:
            if idx < len(valid_stocks):
                top_stocks.add(valid_stocks[idx])
    if len(optimal_weights2) > 0:
        top_indices2 = np.argsort(optimal_weights2)[-3:]
        for idx in top_indices2:
            if idx < len(valid_stocks):
                top_stocks.add(valid_stocks[idx])

    for stock in top_stocks:
        report += f"- {stock}\n"

    report += "\n注意: 实际投资决策应结合市场环境和个人风险偏好进行调整。"

    with open('portfolio_strategy_report.txt', 'w', encoding='utf-8') as f:
        f.write(report)

    print("已生成策略报告: portfolio_strategy_report.txt")

def main():
    """主函数"""
    print("=" * 70)
    print("多支股票组合投资策略")
    print(f"运行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 70)

    # 1. 加载股票数据
    stock_data, valid_stocks = load_stock_data()

    if len(valid_stocks) < 2:
        print("错误: 有效股票数量不足，至少需要2支股票")
        return

    # 2. 预处理数据
    returns_matrix, cov_matrix, valid_stocks = preprocess_data(stock_data, valid_stocks)

    # 检查预处理结果
    if returns_matrix is None or cov_matrix is None or len(valid_stocks) < 2:
        print("错误: 数据预处理失败，无法继续")
        return

    # 3. 调整收益率以适应交易方案
    adjusted_data = adjust_returns_for_trading_scheme(stock_data, valid_stocks, returns_matrix)

    # 4. 优化投资组合
    print("\n" + "=" * 50)
    print("优化投资组合")
    print("=" * 50)

    # 方案1优化
    print("\n优化方案1 (买入=最低价，卖出=最高价)...")
    optimal_weights1, portfolio_return1, portfolio_volatility1, sharpe_ratio1 = optimize_portfolio(
        adjusted_data['scheme1']['annualized_returns'],
        adjusted_data['scheme1']['cov_matrix']
    )

    # 方案2优化
    print("\n优化方案2 (买入=1/4分位，卖出=3/4分位)...")
    optimal_weights2, portfolio_return2, portfolio_volatility2, sharpe_ratio2 = optimize_portfolio(
        adjusted_data['scheme2']['annualized_returns'],
        adjusted_data['scheme2']['cov_matrix']
    )

    # 5. 生成有效前沿
    print("\n" + "=" * 50)
    print("生成有效前沿")
    print("=" * 50)

    print("\n生成方案1有效前沿...")
    efficient_frontier1 = generate_efficient_frontier(
        adjusted_data['scheme1']['annualized_returns'],
        adjusted_data['scheme1']['cov_matrix']
    )

    print("\n生成方案2有效前沿...")
    efficient_frontier2 = generate_efficient_frontier(
        adjusted_data['scheme2']['annualized_returns'],
        adjusted_data['scheme2']['cov_matrix']
    )

    # 6. 回测投资组合
    print("\n" + "=" * 50)
    print("投资组合回测")
    print("=" * 50)

    # 确定回测期（取最近6个月）
    all_dates = []
    for stock in valid_stocks:
        all_dates.extend(stock_data[stock]['日期'].tolist())

    all_dates = sorted(list(set(all_dates)))
    if len(all_dates) == 0:
        print("错误: 没有可用的日期数据用于回测")
        return

    backtest_start = all_dates[-120] if len(all_dates) > 120 else all_dates[0]
    backtest_end = all_dates[-1]

    print(f"\n回测期: {backtest_start.date()} 至 {backtest_end.date()}")

    print("\n回测方案1...")
    backtest_results1 = backtest_portfolio(
        stock_data, optimal_weights1, valid_stocks,
        backtest_start, backtest_end
    )

    print("\n回测方案2...")
    backtest_results2 = backtest_portfolio(
        stock_data, optimal_weights2, valid_stocks,
        backtest_start, backtest_end
    )

    # 7. 分析和可视化结果
    plot_results(
        adjusted_data, valid_stocks,
        efficient_frontier1, efficient_frontier2,
        optimal_weights1, optimal_weights2,
        backtest_results1, backtest_results2
    )

    # 8. 保存结果
    save_results(
        stock_data,
        valid_stocks, optimal_weights1, optimal_weights2,
        efficient_frontier1, efficient_frontier2,
        backtest_results1, backtest_results2,
        adjusted_data
    )

    # 9. 打印最终结论
    print("\n" + "=" * 70)
    print("多支股票组合投资策略完成！")
    print("=" * 70)

    print(f"方案1 (买入=最低价，卖出=最高价):")
    print(f"  最优夏普比率: {efficient_frontier1['tangent']['sharpe']:.4f}")
    print(f"  年化收益率: {efficient_frontier1['tangent']['return']:.4%}")
    print(f"  年化波动率: {efficient_frontier1['tangent']['risk']:.4%}")

    print(f"\n方案2 (买入=1/4分位，卖出=3/4分位):")
    print(f"  最优夏普比率: {efficient_frontier2['tangent']['sharpe']:.4f}")
    print(f"  年化收益率: {efficient_frontier2['tangent']['return']:.4%}")
    print(f"  年化波动率: {efficient_frontier2['tangent']['risk']:.4%}")

    if efficient_frontier1['tangent']['sharpe'] > efficient_frontier2['tangent']['sharpe']:
        print("\n结论: 方案1 (买入=最低价，卖出=最高价) 提供了更高的风险调整后收益")
    else:
        print("\n结论: 方案2 (买入=1/4分位，卖出=3/4分位) 提供了更高的风险调整后收益")

    print("\n结果文件:")
    print("  - stock_risk_return.png: 股票收益率-风险分析")
    print("  - efficient_frontier.png: 有效前沿对比")
    print("  - portfolio_weights.png: 投资组合权重对比")
    print("  - portfolio_value.png: 投资组合价值对比")
    print("  - portfolio_composition.png: 投资组合构成表格")
    print("  - optimal_weights.csv: 最优权重")
    print("  - efficient_frontier_scheme1.csv: 方案1有效前沿数据")
    print("  - efficient_frontier_scheme2.csv: 方案2有效前沿数据")
    print("  - backtest_results_scheme1.csv: 方案1回测结果")
    print("  - backtest_results_scheme2.csv: 方案2回测结果")
    print("  - portfolio_strategy_report.txt: 投资组合策略报告")
    print("=" * 70)


if __name__ == "__main__":
    main()